package main

import (
	"fmt"
	"io/fs"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"unicode/utf8"
)

// DocumentChunk 文档分片
type DocumentChunk struct {
	Content   string
	LineStart int
	LineEnd   int
	Index     int
}

// DocumentProcessor 文档处理器
type DocumentProcessor struct {
	chunkSize int
}

// NewDocumentProcessor 创建文档处理器
func NewDocumentProcessor(chunkSize int) *DocumentProcessor {
	return &DocumentProcessor{
		chunkSize: chunkSize,
	}
}

// ProcessFolder 处理文件夹
func (dp *DocumentProcessor) ProcessFolder(folderPath string, fileTypes []string) ([]string, error) {
	var files []string
	
	// 创建文件类型映射
	typeMap := make(map[string]bool)
	for _, ext := range fileTypes {
		if !strings.HasPrefix(ext, ".") {
			ext = "." + ext
		}
		typeMap[strings.ToLower(ext)] = true
	}
	
	err := filepath.WalkDir(folderPath, func(path string, d fs.DirEntry, err error) error {
		if err != nil {
			fmt.Printf("❌ 访问文件失败 %s: %v\n", path, err)
			return nil // 继续处理其他文件
		}
		
		if d.IsDir() {
			return nil
		}
		
		ext := strings.ToLower(filepath.Ext(path))
		if typeMap[ext] {
			files = append(files, path)
			fmt.Printf("📄 发现文件: %s\n", path)
		}
		
		return nil
	})
	
	if err != nil {
		return nil, fmt.Errorf("遍历文件夹失败: %v", err)
	}
	
	fmt.Printf("📊 共发现 %d 个匹配文件\n", len(files))
	return files, nil
}

// ProcessFile 处理单个文件
func (dp *DocumentProcessor) ProcessFile(filePath string) ([]DocumentChunk, error) {
	fmt.Printf("🔄 正在处理文件: %s\n", filePath)
	
	// 读取文件内容
	content, err := dp.readFileWithEncoding(filePath)
	if err != nil {
		return nil, fmt.Errorf("读取文件失败: %v", err)
	}
	
	// 分割成行
	lines := strings.Split(content, "\n")
	fmt.Printf("📝 文件共 %d 行\n", len(lines))
	
	// 进行分片
	chunks := dp.chunkText(lines)
	fmt.Printf("✂️  分片完成，共 %d 个分片\n", len(chunks))
	
	return chunks, nil
}

// readFileWithEncoding 读取文件并处理编码
func (dp *DocumentProcessor) readFileWithEncoding(filePath string) (string, error) {
	data, err := os.ReadFile(filePath)
	if err != nil {
		return "", err
	}
	
	// 检查是否为有效的UTF-8
	if utf8.Valid(data) {
		return string(data), nil
	}
	
	// 尝试其他编码（简化处理，实际项目中可以使用更复杂的编码检测）
	fmt.Printf("⚠️  文件可能不是UTF-8编码，尝试强制转换: %s\n", filePath)
	return string(data), nil
}

// chunkText 文本分片
func (dp *DocumentProcessor) chunkText(lines []string) []DocumentChunk {
	var chunks []DocumentChunk
	var currentChunk strings.Builder
	var chunkStartLine int
	chunkIndex := 0
	
	for i, line := range lines {
		// 如果当前分片为空，记录起始行
		if currentChunk.Len() == 0 {
			chunkStartLine = i + 1
		}
		
		// 添加行到当前分片
		if currentChunk.Len() > 0 {
			currentChunk.WriteString("\n")
		}
		currentChunk.WriteString(line)
		
		// 检查是否需要分片
		shouldSplit := false
		
		// 按大小分片
		if currentChunk.Len() >= dp.chunkSize {
			shouldSplit = true
		}
		
		// 按语义分片（段落结束）
		if strings.TrimSpace(line) == "" && currentChunk.Len() > 100 {
			shouldSplit = true
		}
		
		// 按标题分片（Markdown标题）
		if strings.HasPrefix(strings.TrimSpace(line), "#") && currentChunk.Len() > 100 {
			// 将标题放到下一个分片
			content := currentChunk.String()
			lastNewline := strings.LastIndex(content, "\n")
			if lastNewline > 0 {
				// 创建当前分片（不包含标题）
				chunks = append(chunks, DocumentChunk{
					Content:   content[:lastNewline],
					LineStart: chunkStartLine,
					LineEnd:   i,
					Index:     chunkIndex,
				})
				
				// 开始新分片（包含标题）
				currentChunk.Reset()
				currentChunk.WriteString(line)
				chunkStartLine = i + 1
				chunkIndex++
				continue
			}
		}
		
		if shouldSplit {
			content := strings.TrimSpace(currentChunk.String())
			if content != "" {
				chunks = append(chunks, DocumentChunk{
					Content:   content,
					LineStart: chunkStartLine,
					LineEnd:   i + 1,
					Index:     chunkIndex,
				})
				chunkIndex++
			}
			currentChunk.Reset()
		}
	}
	
	// 处理最后一个分片
	if currentChunk.Len() > 0 {
		content := strings.TrimSpace(currentChunk.String())
		if content != "" {
			chunks = append(chunks, DocumentChunk{
				Content:   content,
				LineStart: chunkStartLine,
				LineEnd:   len(lines),
				Index:     chunkIndex,
			})
		}
	}
	
	return chunks
}

// ExtractSummary 提取摘要
func (dp *DocumentProcessor) ExtractSummary(content string) string {
	// 简单的摘要提取：取前200个字符
	runes := []rune(content)
	if len(runes) <= 200 {
		return content
	}
	
	summary := string(runes[:200])
	
	// 尝试在句号处截断
	if idx := strings.LastIndex(summary, "。"); idx > 100 {
		summary = summary[:idx+3] // 包含句号
	} else if idx := strings.LastIndex(summary, "."); idx > 100 {
		summary = summary[:idx+1]
	}
	
	return summary + "..."
}

// ExtractKeywords 提取关键词
func (dp *DocumentProcessor) ExtractKeywords(content string) []string {
	// 简单的关键词提取：基于词频
	words := dp.segmentText(content)
	wordCount := make(map[string]int)
	
	for _, word := range words {
		if len(word) >= 2 { // 只考虑长度>=2的词
			wordCount[word]++
		}
	}
	
	// 按频率排序
	type wordFreq struct {
		word  string
		count int
	}
	
	var frequencies []wordFreq
	for word, count := range wordCount {
		if count >= 2 { // 至少出现2次
			frequencies = append(frequencies, wordFreq{word, count})
		}
	}
	
	// 简单排序
	for i := 0; i < len(frequencies)-1; i++ {
		for j := i + 1; j < len(frequencies); j++ {
			if frequencies[i].count < frequencies[j].count {
				frequencies[i], frequencies[j] = frequencies[j], frequencies[i]
			}
		}
	}
	
	// 返回前10个关键词
	var keywords []string
	limit := 10
	if len(frequencies) < limit {
		limit = len(frequencies)
	}
	
	for i := 0; i < limit; i++ {
		keywords = append(keywords, frequencies[i].word)
	}
	
	return keywords
}

// segmentText 简单的中文分词
func (dp *DocumentProcessor) segmentText(text string) []string {
	// 移除标点符号和特殊字符
	reg := regexp.MustCompile(`[^\p{L}\p{N}]+`)
	cleaned := reg.ReplaceAllString(text, " ")
	
	// 按空格分割
	words := strings.Fields(cleaned)
	
	// 对于中文，按字符分割（简化处理）
	var result []string
	for _, word := range words {
		if len(word) > 0 {
			// 检查是否包含中文字符
			if regexp.MustCompile(`[\p{Han}]`).MatchString(word) {
				// 中文按2-3字符组合
				runes := []rune(word)
				for i := 0; i < len(runes); i++ {
					if i+1 < len(runes) {
						result = append(result, string(runes[i:i+2]))
					}
					if i+2 < len(runes) {
						result = append(result, string(runes[i:i+3]))
					}
				}
			} else {
				// 英文直接添加
				result = append(result, strings.ToLower(word))
			}
		}
	}
	
	return result
}